from typing import List, Tuple
from pathlib import Path
from PIL import Image
import io
from .gemini_manager import GeminiManager
from ..config import settings


class ImageGenerator:
    """图片生成器"""
    
    def __init__(self, gemini_manager: GeminiManager):
        self.gemini = gemini_manager
    
    async def generate_batch(self, prompts: List[str]) -> List[Tuple[str, Path]]:
        """批量生成图片
        
        Args:
            prompts: 图片生成提示词列表
            
        Returns:
            List[Tuple[str, Path]]: 提示词和对应生成图片路径的列表
        """
        results = []
        for i, prompt in enumerate(prompts):
            try:
                image_data = await self.generate_image(prompt)
                if self.validate_image(image_data):
                    path = self._save_image(image_data, i)
                    results.append((prompt, path))
            except Exception as e:
                print(f"生成图片失败 {i}: {str(e)}")
                continue
        return results
    
    async def generate_image(self, prompt: str) -> bytes:
        """生成单张图片"""
        return await self.gemini.generate_image(prompt)
    
    def validate_image(self, image_data: bytes) -> bool:
        """验证图片质量"""
        try:
            image = Image.open(io.BytesIO(image_data))
            # TODO: 实现图片质量验证逻辑
            return True
        except Exception:
            return False
    
    def _save_image(self, image_data: bytes, index: int) -> Path:
        """保存图片到本地"""
        output_dir = Path(settings.IMAGE_OUTPUT_PATH)
        output_dir.mkdir(parents=True, exist_ok=True)
        
        image_path = output_dir / f"wallpaper_{index:03d}.jpg"
        with open(image_path, "wb") as f:
            f.write(image_data)
        
        return image_path
